import os
os.environ['CUDA_VISIBLE_DEVICES'] = '0'

from swift.llm import (
    get_model_tokenizer, get_template, inference, ModelType, get_default_template_type
)
from swift.tuners import Swift

ckpt_dir = '../tuned_llms/qwen-1_8b-chat-int4/v6-20240327-160204/checkpoint-56'  # tuned weights (not including the original weights)
model_type = ModelType.qwen1half_1_8b_chat_int4
template_type = get_default_template_type(model_type)
print(template_type)
model_dir = "../pretrain_llms/qwen/Qwen-1_8B-Chat-Int4"  # original weights
model, tokenizer = get_model_tokenizer(model_type, model_dir=model_dir, model_kwargs={'device_map': 'auto'})

model = Swift.from_pretrained(model, ckpt_dir, inference_mode=True)
template = get_template(template_type, tokenizer)
print(template)
query = '向右前方45度飞行10米'
response, history = inference(model, template, query,system='''你是一名助手，帮助我使用无人机的AirSim模拟器。
当我要求你做某事时，你应该给我使用AirSim完成任务所需的Python代码，然后解释该代码的作用。
您只能使用我为您定义的函数。
你不能使用任何其他你认为可能存在的假设函数。
您可以使用math和numpy等库中的简单Python函数。

以下我定义的函数，是一些可以用来指挥无人机的功能：
aw.takeoff() - 起飞无人机。
aw.land() - 降落无人机。
aw.get_drone_position() - 将无人机的当前位置返回为与X、Y、Z坐标相对应的3个浮点数。
aw.fly_to([X, Y, Z]) - 将无人机飞行到指定的位置，该位置由对应于X、Y、Z坐标的三个参数组成。
aw.fly_to([0, 0, 0]) - 将无人机回到原点位置
aw.fly_path(points) - 使无人机沿着点数组指定的路径飞行。每个点也是一个由3个浮点数组成的数组，这些浮点数对应于X、Y、Z坐标。
aw.set_yaw(yaw) - 将无人机的偏航设置为以度为单位的指定值。
aw.get_yaw() - 以度为单位返回无人机的当前偏航。
aw.get_position(object_name): 以一个字符串作为输入，指示感兴趣对象的名称，并返回一个由3个浮点数组成的数组，指示其X、Y、Z坐标。

根据轴惯例，正向表示正X轴。右侧表示正Y轴。向上表示正Z轴。
原点坐标是[0,0,0]
无人机向左移动：Y轴减少数值
无人机向右移动：Y轴增加数值
无人机向上移动：Z轴增加数值
无人机向下移动：Z轴减少数值
无人机向后移动：X轴减少数值
无人机向前移动：X轴增加数值''')
print(f'response: {response}')
print(f'history: {history}')
